# OPLogin.py文件
import re, json, jsonlines
from locust import HttpUser,task,TaskSet,between, events,User
import asyncio
from httpx import AsyncClient
import requests
import websocket, time



class MyTask(TaskSet):
    def on_start(self):
        self.params = {
            'messages':[{
                'role': 'user',
                'content': '你是谁'
            }],
            'stream': True
        }
        self.headers = {"Content-Type": "application/json"}
        print("********** 测试开始 **********")
   
    def on_stop(self):
        print("********** 测试结束 **********")


    @task
    def task_chat(self):
        url = '/chat/local/chat'
        response = self.client.post(url, headers=self.headers, json=self.params)
        try:
            assert response.status_code == 200
            response = json.loads(response.text)
            response = response.get('response').get('content')
            print(f'from {url} : {response}')
        except Exception as e:
            print(f"http error: code {response.status_code} | {e}")
         
    # @task
    # def task_websocket_chat(self):
    #     url = 'ws://127.0.0.1:8000/chat/local/websocket_chat'
    #     try:
    #         socket = websocket.create_connection(url, timeout=1000)
    #         socket.send(json.dumps(self.params))
    #         response_collect = ''
    #         while True:
    #             response = socket.recv()
    #             response = json.loads(response)
    #             if response.get('finish'):
    #                 websocket.close()
    #                 break
    #             response_collect += response.get('response').get('content')
    #         print(response_collect)
    #     except Exception as e:
    #         print('error at: ',e)
        
    @task 
    def task_stream_chat(self):
        url = '/chat/local/stream_chat'
        response = self.client.post(url, headers=self.headers, json=self.params)
        try:
            assert response.status_code == 200
            response = response.text
            response = response.replace('data: ', '').split('\r\n\r\n')
            response = ''.join(response)
            print(f'from {url} : {response}')
        except Exception as e:
            print(f"http error: code {response.status_code} | {e}")




class OPlogin(HttpUser):
    tasks = [MyTask]				# 要执行的任务
    host = "http://127.0.0.1:8000"   # 被测网站域名
    wait_time = between(1,3)			# 请求间隔时间
